import numpy as np
import os
import re
from os import listdir
from os.path import join
from scipy import io
import pandas as pd
import struct
from torch.utils.data import DataLoader, Dataset
from process import *

mapping = {
    'climb': 0,
    'stand': 1,
    'jump': 2,
    'fall': 3,
    'sit': 4,
    'get up': 5,
    'walk': 6,
    'run': 7,
    'lift': 8,
    'lie': 9,
    'bend': 10,
    'pick': 11
}

# 事件流处理方法
repr_map = {'eventFrame': get_eventFrame,
            'eventAccuFrame': get_eventAccuFrame,
            'timeSurface': get_timeSurface,
            'eventCount': get_eventCount}


# left or right move all event locations randomly
def random_shift_events(events, max_shift=20, resolution=(180, 240)):
    """
    left or right move all event locations randomly
    event locations: (t, x, y, p)
    """
    H, W = resolution
    x_shift, y_shift = np.random.randint(-max_shift, max_shift + 1, size=(2,))
    events[:, 1] += x_shift
    events[:, 2] += y_shift

    valid_events = (events[:, 1] >= 0) & (events[:, 1] < W) & (events[:, 2] >= 0) & (events[:, 2] < H)
    events = events[valid_events]

    return events


# flip half of the event images along the x dimension
def random_flip_events_along_x(events, resolution=(180, 240), p=0.5):
    """
    event locations: (t, x, y, p)
    """
    H, W = resolution
    if np.random.random() < p:
        events[:, 1] = W - 1 - events[:, 1]
    return events


class DailyActionDataset(Dataset):
    def __init__(self, root, train=True, augmentation=False,
                 repr=['timeSurface'], time_num=20, frame_transform=None, label_transform=None):
        """
        T:T,即为时间窗口数量(即为将事件流的分成T/chunk块)
        repr:事件流处理方法
        repr_map = {'eventFrame':get_eventFrame,
                    'eventAccuFrame':get_eventAccuFrame,
                    'timeSurface':get_timeSurface,
                    'eventCount':get_eventCount}
        """
        super(DailyActionDataset, self).__init__()
        self.train = train
        self.files = []
        self.labels = []
        self.augmentation = augmentation
        self.root = root
        self.frame_transform = frame_transform
        self.label_transform = label_transform
        self.repr = repr
        self.time_num = time_num

        if train:
            self.class_dir_train = os.listdir(self.root)
        else:
            self.class_dir_test = os.listdir(self.root)

        if train:
            root_train = self.root
            self.files_train = []
            for self.now_file, self.label_train_one, self.files_train_one in os.walk(root_train):
                i = 0
                for now_data_file in self.files_train_one:
                    if i < 120 * 0.8:
                        self.files_train.append(os.path.join(self.now_file, now_data_file))
                    i += 1
            self.n = len(self.files_train)
        else:
            root_test = self.root
            self.files_test = []
            for self.now_file, self.label_train_one, self.files_train_one in os.walk(root_test):
                i = 0
                for now_data_file in self.files_train_one:
                    if i >= 120 * 0.8:
                        self.files_test.append(os.path.join(self.now_file, now_data_file))
                    i += 1
            self.n = len(self.files_test)

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        if self.train:
            file_path = self.files_train[idx]
        else:
            file_path = self.files_test[idx]
        label = mapping[file_path.split('\\')[-2]]
        data_dvs = np.array(getDVSeventsDavis(file_path), dtype=np.int64)
        # normalize the timestamps
        _min = data_dvs[0, :].min()
        _max = data_dvs[0, :].max()
        data_dvs[0, :] = (data_dvs[0, :] - _min) / (_max - _min)

        # 将(4,事件数量)转置为(事件数量,4)
        data_dvs = data_dvs.transpose()

        if self.augmentation:
            data_dvs = random_shift_events(data_dvs)
            data_dvs = random_flip_events_along_x(data_dvs)

        reprs = []
        for repr_name in self.repr:
            # 输入(t,x,y,p), 输出(T,y,x) y:[0,260] x:[0:346]
            repr_array = repr_map[repr_name](data_dvs[:, 0], data_dvs[:, 1].astype(np.int32),
                                             data_dvs[:, 2].astype(np.int32),
                                             data_dvs[:, 3],
                                             repr_size=(260, 346), time_num=self.time_num)

            # standardization
            mu = np.mean(repr_array)
            sigma = np.std(repr_array)
            repr_array = (repr_array - mu) / sigma

            reprs.append(repr_array)

        reprs = np.array(reprs)  # (C,T,260,346)
        reprs = np.transpose(reprs, (1, 0, 3, 2))  # (T,C,346,260) (y,x)->(x,y)

        if self.frame_transform is not None:
            reprs = self.frame_transform(reprs)

        if self.label_transform is not None:
            label = self.label_transform(label)

        return reprs, label


def getDVSeventsDavis(file, numEvents=1e10, startTime=0):
    """ DESCRIPTION: This function reads a given aedat file and converts it into four lists indicating
                     timestamps, x-coordinates, y-coordinates and polarities of the event stream.

    Args:
        file: the path of the file to be read, including extension (str).
        numEvents: the maximum number of events allowed to be read (int, default value=1e10).
        startTime: the start event timestamp (in microseconds) where the conversion process begins (int, default value=0).

    Return:
        ts: list of timestamps in microseconds.
        x: list of x-coordinates in pixels.
        y: list of y-coordinates in pixels.`
        pol: list of polarities (0: on -> off, 1: off -> on).
    """
    # print('\ngetDVSeventsDavis function called \n')
    sizeX = 346
    sizeY = 260
    x0 = 0
    y0 = 0
    x1 = sizeX
    y1 = sizeY

    # print('Reading in at most', str(numEvents))

    triggerevent = int('400', 16)
    polmask = int('800', 16)
    xmask = int('003FF000', 16)
    ymask = int('7FC00000', 16)
    typemask = int('80000000', 16)
    typedvs = int('00', 16)
    xshift = 12
    yshift = 22
    polshift = 11
    x = []
    y = []
    ts = []
    pol = []
    numeventsread = 0

    length = 0
    aerdatafh = open(file, 'rb')
    k = 0
    p = 0
    statinfo = os.stat(file)
    if length == 0:
        length = statinfo.st_size
    # print("file size", length)

    lt = aerdatafh.readline()
    while lt and str(lt)[2] == "#":
        p += len(lt)
        k += 1
        lt = aerdatafh.readline()
        continue

    aerdatafh.seek(p)
    tmp = aerdatafh.read(8)
    p += 8
    while p < length:
        ad, tm = struct.unpack_from('>II', tmp)
        ad = abs(ad)
        if tm >= startTime:
            if (ad & typemask) == typedvs:
                xo = sizeX - 1 - float((ad & xmask) >> xshift)
                yo = float((ad & ymask) >> yshift)
                polo = 1 - float((ad & polmask) >> polshift)
                if xo >= x0 and xo < x1 and yo >= y0 and yo < y1:
                    x.append(xo)
                    y.append(yo)
                    pol.append(polo)
                    ts.append(tm)
        aerdatafh.seek(p)
        tmp = aerdatafh.read(8)
        p += 8
        numeventsread += 1

    # print('Total number of events read =', numeventsread)
    # print('Total number of DVS events returned =', len(ts))

    return ts, x, y, pol


